#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""Image recognition with PCA

Ref.
M. A. Turk A. P. Pentland. Face Recognition Using Eigenfaces, 1991.

Example:
    images = a list of image (read by pillow)
    ip = ImagePCA(n_components=2)
    ip.fit(images) # or ip.ezfit(folder) where folder is the folder of images
    # get and save eigen images
    for k, im in enumerate(ip.eigen_images):
        im.save(f'eigen{k}.jpg')
    # get coordinates and errors of images (that could be new)
    coords = ip.coordinate(images)
    errors = ip.error(images)

    # see demo function
"""


import numpy as np
import numpy.linalg as LA

from sklearn.decomposition import *
from sklearn.base import TransformerMixin
import PIL_ext

def op_images(m):
    def mm(obj, images, *args, **kwargs):
        size = images[0].size
        mode = images[0].mode
        X = PIL_ext.tomatrix(images, 'row', size)
        if obj.size is None:
            obj.size = size
        if obj.mode is None:
            obj.mode = mode
        return m(obj, X, *args, **kwargs)
    return mm


class ImagePCA(PCA):
    """PCA Transformer for images

    eigens_: eigen images in vector form, KL basis multiplied by singular values
    """

    def __init__(self, size=None, mode=None, *args, **kwargs):
        """
        Keyword Arguments:
            size {[type]} -- size of image (default: {None})
            mode {[type]} -- mode of image (default: {None})
        """
        super(ImagePCA, self).__init__(*args, **kwargs)
        self.size = size
        self.mode = mode

    @op_images
    def fit(self, X):
        """fit method
        
        Decorators:
            op_images
        
        Arguments:
            X {np.ndarray} -- matrix of images
        """
        super(ImagePCA, self).fit(X)
        self.eigens_ = np.diag(self.singular_values_) @ self.components_
        self.__Xtrain = X
        return self

    @property
    def coords_(self):
        return super(ImagePCA, self).transform(self.__Xtrain) @ np.diag(1/self.singular_values_)

    @op_images
    def transform(self, X):
        return super(ImagePCA, self).transform(X)

    def fit_transform(self, X):
        return self.fit(X).transform(X)

    @property
    def eigen_images(self):
        """Get eigen images

        transform the eigens_ to images
        """
        return [PIL_ext.toimage(y, self.size, mode=self.mode) for y in self.eigens_]
    
    def inverse_transform(self, Y, toimage=False):
        # transform to images if toimage=True
        X = super(ImagePCA, self).inverse_transform(Y)
        if toimage:
            return [PIL_ext.toimage(x, size=self.size, mode=self.mode) for x in X]
        else:
            return X

    def coordinate(self, X):
        # coordinate under eigen images
        return self.transform(X) @ np.diag(1/self.singular_values_)

    def reconstruct(self, X, toimage=True):
        # reconstruction of images under KL basis
        return self.inverse_transform(self.transform(X), toimage=toimage)

    def error(self, X):
        # error of reconstruction
        R = self.reconstruct(X, toimage=False)
        X = PIL_ext.tomatrix(X, 'row', size=self.size)
        return LA.norm(X-R, axis=1)

    def ezfit(self, folder='images'):
        images = PIL_ext.get_images(folder=folder)
        self.fit(images)

    def dist(self, X):
        from scipy.spatial.distance import cdist
        return cdist(self.coords_, self.coordinate(X))

    def classify(X):
        D = self.dist(X)
        return D.argmin(axis=0)

    def locate(image):
        pass

    def generate(self, size=1, toimage=False):
        from scipy.stats import gaussian_kde
        cc = super(ImagePCA, self).transform(self.__Xtrain).T
        c = np.row_stack([gaussian_kde(c).resample(size) for c in cc])
        return self.inverse_transform(c.T, toimage=toimage)


def demo(*args, **kwargs):
    # save images in 'images' folder before pca
    ip = ImagePCA(*args, **kwargs)
    ip.ezfit(folder='images')
    for k, im in enumerate(ip.eigen_images):
        im.save(f'eigen{k}.jpg')
    for k, im in enumerate(ip.generate(10, toimage=True)):
        im.save(f'artificial{k}.jpg')

    # output coordinates of images under eigen images
    # coords = ip.coordinate(images)
    # errors = ip.error(images)
    # import pandas as pd
    # df = pd.DataFrame(data=np.column_stack([coords, errors]), columns=tuple(f'eigen{k}.jpg' for k in range(coords.shape[1]))+('error',))
    # data = df.to_csv('test.csv')

if __name__ == '__main__':
    demo(n_components=5)
    